package com.dtdream.emr;

import org.apache.hadoop.hive.ql.exec.Description;
import org.apache.hadoop.hive.ql.exec.UDFArgumentException;
import org.apache.hadoop.hive.ql.metadata.HiveException;
import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.ObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.PrimitiveObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.StructObjectInspector;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory;
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorUtils;

import java.util.Arrays;
import java.util.List;

/**
 * Created by moriarty on 17-8-4.
 */
@Description("_FUNC_(separator, col) - split column into several rows by separator")
public class GenericUDTFSplit extends GenericUDTF {

    private PrimitiveObjectInspector separatorOI;
    private PrimitiveObjectInspector valueOI;

    @Override
    public StructObjectInspector initialize(ObjectInspector[] arguments) throws UDFArgumentException {
        if (arguments.length != 2) {
            throw new UDFArgumentException("Function GenericUDTFSplit requires 2 arguments, got " + arguments.length);
        }

        valueOI = (PrimitiveObjectInspector) arguments[0];
        separatorOI = (PrimitiveObjectInspector) arguments[1];


        if (separatorOI.getPrimitiveCategory() != PrimitiveObjectInspector.PrimitiveCategory.STRING
                || valueOI.getPrimitiveCategory() != PrimitiveObjectInspector.PrimitiveCategory.STRING) {
            throw new UDFArgumentException("Both separator and value should be string.");
        }

        List<String> fieldNames = Arrays.asList("col1");
        List<ObjectInspector> fieldOIs = Arrays.asList((ObjectInspector) PrimitiveObjectInspectorFactory.javaStringObjectInspector);
        return ObjectInspectorFactory.getStandardStructObjectInspector(fieldNames, fieldOIs);
    }

    @Override
    public void process(Object[] args) throws HiveException {
        String value = PrimitiveObjectInspectorUtils.getString(args[0], valueOI);
        String separator = PrimitiveObjectInspectorUtils.getString(args[1], separatorOI);


        if (value != null) {
            for (String substr : value.split(separator)) {
                forward(Arrays.asList(substr));
            }
        }
    }

    @Override
    public void close() throws HiveException {
    }
}
